use std::pin::Pin;
use std::sync::Arc;
use async_trait::async_trait;
use futures::Stream;
use tokio::net::TcpStream;
use tokio::net::tcp::OwnedWriteHalf;
use tokio::sync::{Mutex, broadcast};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use super::{McpMessage, McpResult, McpTransport};
use crate::McpError;

pub struct TcpTransport {
    writer: Arc<Mutex<OwnedWriteHalf>>,
    receiver: broadcast::Receiver<McpResult<McpMessage>>,
}

impl TcpTransport {
    pub async fn new(stream: TcpStream) -> (Self, broadcast::Sender<McpResult<McpMessage>>) {
        // 将一个 TcpStream 分为 Read 和 Write 端
        let (reader, writer) = stream.into_split();
        // 将 write 端口封装到 Arc<Mutex> 中
        let writer = Arc::new(Mutex::new(writer));
        // 创建一个异步通道（有 send 和 receive 端）
        let (sender, receiver) = broadcast::channel(100);

        // 给通道读取端口增加 buf
        let mut reader = BufReader::new(reader);
        let sender_clone = sender.clone();

        // 后台读取任务
        tokio::spawn(async move {
            let mut line = String::new();
            loop {
                line.clear();
                // 是要 TcpStream 获得数据
                match reader.read_line(&mut line).await {
                    Ok(0) => break, // 连接关闭
                    Ok(_) => {
                        // 视为 json 格式解析为 McpMessage！
                        let message = match serde_json::from_str(&line) {
                            Ok(msg) => Ok(msg),
                            Err(err) => Err(McpError::Serialization(err.to_string())),
                        };
                        if sender_clone.send(message).is_err() {
                            break; // 如果没人接收，退出
                        }
                    }
                    Err(err) => {
                        let _ = sender_clone.send(Err(McpError::Io(err.to_string())));
                        break;
                    }
                }
            }
        });

        (
            Self { writer, receiver },
            sender
        )
    }
}

#[async_trait]
impl McpTransport for TcpTransport {
    async fn send(&self, message: McpMessage) -> McpResult<()> {
        // 获取 Tcp 的输出端口
        let mut writer = self.writer.lock().await;
        let json = serde_json::to_string(&message)?;
        writer.write_all(json.as_bytes()).await.map_err(|e| McpError::Io(e.to_string()))?;
        writer.write_all(b"\n").await.map_err(|e| McpError::Io(e.to_string()))?;
        writer.flush().await.map_err(|e| McpError::Io(e.to_string()))?;
        Ok(())
    }

    fn receive(&self) -> Pin<Box<dyn Stream<Item = McpResult<McpMessage>> + Send>> {
        let rx = self.receiver.resubscribe();
        Box::pin(futures::stream::unfold(rx, |mut rx| async move {
            match rx.recv().await {
                Ok(msg) => Some((msg, rx)),
                Err(_) => None,
            }
        }))
    }

    async fn close(&self) -> McpResult<()> {
        // 实际上 TcpStream 的 close 是 Drop 自动完成的
        Ok(())
    }
}
